import yaml
import os
import argparse

yamlpath = './configs/electronic_eye/_base_/tta_reader.yml'

multi_scale_list = [
    {'Resize': {'interp':2, 'target_size': [1080,1920], 'keep_ratio': True}},
    {'Resize': {'interp':2, 'target_size': [900, 1600], 'keep_ratio': True}},
    {'Resize': {'interp':2, 'target_size': [810, 1440], 'keep_ratio': True}},
    {'Resize': {'interp':2, 'target_size': [720, 1280], 'keep_ratio': True}},
]

sample_transforms_pipeline = [
    {'Decode': {}},
    {'Resize': {'interp':2, 'target_size': [1080,1920], 'keep_ratio': True}},
    #{'TTA': {}},
    {'NormalizeImage': {'is_scale': True, 'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}},
    {'Permute': {}}
]



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument(
    "--yamlpath",
    default=r'D:\Soft\python\machine_learning\cud\PaddleDetection-release-2.0-rc\dygraph\configs\cascade_rcnn\_base_\tta_reader.yml',
    type=str,
    help="your yaml path"
    )

    parser.add_argument(
    "--sizeid",
    default=1,
    type=int,
    help="resize id: 0-3"
    )

    parser.add_argument(
        '--only_resize',
        action='store_true',
        default=True,
        help='is only resize')

    args = parser.parse_args()

    assert args.sizeid <= 3
    print("change yml sizeid", args.sizeid)
    with open(args.yamlpath, 'r', encoding='utf-8') as f:
        data = f.read()
        yaml_data = yaml.load(data, Loader=yaml.FullLoader)

    sample_transforms_pipeline[1] = multi_scale_list[args.sizeid]

    yaml_data['EvalReader']['sample_transforms'] = sample_transforms_pipeline

    with open(args.yamlpath, 'w', encoding='utf-8') as f:
        yaml.dump(yaml_data, f)